#include <stdio.h>
#include <stdlib.h>

#define COLOR_BLACK 'b'
#define COLOR_RED   'r'

typedef struct rbtree_node_s rbtree_node_t; 
typedef struct rbtree_node_s* rbtree_t; 

struct rbtree_node_s
{
    int  data;
    char color;
    struct rbtree_node_s *parent;
    struct rbtree_node_s *lchild;
    struct rbtree_node_s *rchild;
};

void set_color_black(rbtree_node_t *node)
{
    node->color = COLOR_BLACK;
}

void set_color_red(rbtree_node_t *node)
{
    node->color = COLOR_RED;
}

int is_color_black(rbtree_node_t *node)
{
    if (node == NULL)
        return 1;
    return node->color == COLOR_BLACK ? 1:0;
}

int is_color_red(rbtree_node_t *node)
{
    if (node == NULL)
        return 0;
    return node->color == COLOR_RED ?1:0;
}

int is_left_child(rbtree_node_t *c)
{
    if (c->parent != NULL)
    {
        return c == c->parent->lchild?1:0;
    }
    return 0;
}

int is_right_child(rbtree_node_t *c)
{
    if (c->parent != NULL)
    {
        return c == c->parent->rchild?1:0;
    }
    return 0;
}

rbtree_node_t *new_rbtree_node(int data)
{
    rbtree_node_t *node = (rbtree_node_t *)malloc(sizeof(rbtree_node_t));
    node->data = data;
    node->parent = NULL;
    node->lchild = node->rchild = NULL;
    set_color_red(node);
}

void rbtree_node_free(rbtree_node_t *n)
{
    if(n->parent != NULL)
        n->parent = NULL;
    if(n->lchild != NULL)
        n->lchild = NULL;
    if(n->rchild != NULL)
        n->rchild = NULL;
    free(n);
}

rbtree_node_t *sibling(rbtree_node_t *c)
{
    if (c->parent == NULL)
        return NULL;
    if (c == c->parent->lchild) 
        return c->parent->rchild;
    else 
        return c->parent->lchild;
}

void left_rotate(rbtree_node_t *c)
{
    rbtree_node_t *right = c->rchild;
    right->parent = c->parent;
    if (c == c->parent->lchild) 
        c->parent->lchild = right;
    else 
        c->parent->rchild = right;
    c->parent = right;
    right->lchild = c;
    c->rchild = NULL;
}

void right_rotate(rbtree_node_t *c)
{
    rbtree_node_t *left = c->lchild;
    left->parent = c->parent;
    if (c == c->parent->lchild)
        c->parent->lchild = left;
    else 
        c->parent->rchild = left;
    c->parent = left;
    left->rchild = c;
    c->lchild = NULL;
}

void rebalance(rbtree_t *root, rbtree_node_t *c)
{
    if (c == *root)
    {
        set_color_black(c);
        *root = c;
        return ;
    }
    rbtree_node_t *parent = c->parent;
    rbtree_node_t *grandpa = parent->parent;
    if (grandpa == NULL)
    {
        return ;
    }
    rbtree_node_t *uncle = sibling(parent);

    if (is_color_red(parent))
    {
        if (is_color_red(uncle))
        {
            set_color_black(parent);
            set_color_black(uncle);
            set_color_red(grandpa);
            c = grandpa;
            rebalance(root, c);
        } 
        else 
        {
            if (is_left_child(c))
            {
                if(is_left_child(parent))
                {
                    set_color_black(parent);
                    set_color_red(grandpa);
                    right_rotate(grandpa);
                }
                else
                {
                    left_rotate(parent);
                    set_color_black(parent);
                    set_color_red(grandpa);
                    right_rotate(grandpa);
                }
            }
            if (is_right_child(c))
            {
                if(is_right_child(parent))
                {
                    set_color_black(parent);
                    set_color_red(grandpa);
                    left_rotate(grandpa);
                }
                else
                {
                    right_rotate(parent);
                    set_color_black(parent);
                    set_color_red(grandpa);
                    left_rotate(grandpa);
                }
            }
        }
    }

}

rbtree_node_t *get_left_max(rbtree_node_t *c)
{
    rbtree_node_t *t = c->lchild;
    while(t != NULL)
    {
        if(t->rchild == NULL)
            break;
        t = t->rchild;
    }
    return t;
}

int rbtree_insert(rbtree_t *root, int data)
{
    rbtree_node_t *p = new_rbtree_node(data);
    if (*root == NULL)
    {
        set_color_black(p);
        *root = p;
        return 0;
    }

    rbtree_node_t *t = *root;
    while(t != NULL)
    {
        if (data < t->data)
        {
            if (t->lchild == NULL)
            {
                t->lchild = p;
                p->parent = t;
                break;
            }
            t = t->lchild;
        }
        else if (data > t->data)
        {
            if (t->rchild == NULL)
            {
                t->rchild = p;
                p->parent = t;
                break;
            }
            t = t->rchild;
        }
        else 
        {
            printf("data %d exist.\n", data);
            return 1;
        }
    }
   
    rebalance(root, p);

    return 0;
}

int rbtree_delete(rbtree_t *root, int data)
{
    rbtree_node_t *t = *root;
    while(t != NULL)
    {
        if (data == t->data)
            break;
        else if (data < t->data)
            t = t->lchild;
        else
            t = t->rchild;
    }

    if (t == NULL)
    {
        return 1;
    }
    // delete and rebalance

    rbtree_node_t *lm = get_left_max(t);
    if (lm == NULL)
    {
        if (t->rchild == NULL)
        {
            if(is_left_child(t))
            {
                t->parent->lchild = NULL;
                rbtree_node_free(t);
                return 0;
            }
            t->parent->rchild = NULL;
            rbtree_node_free(t);
            return 0;
        }
        // TODO 
        if(is_right_child(t))
        {
            t->parent->rchild = t->rchild;
            t->rchild->parent = t->parent;
            rbtree_node_free(t);
            return 0;
        }
        else 
        {
            t->parent->lchild = t->rchild;
            t->rchild->parent = t->parent;
            rbtree_node_free(t);
            return 0;
        }
    }
    if (is_left_child(lm))
        lm->parent->lchild = NULL;
    else 
        lm->parent->rchild = NULL;
    lm->parent = t->parent;
    if (is_left_child(t))
        t->parent->lchild = lm;
    else 
        t->parent->rchild = lm;
    lm->lchild = t->lchild;
    lm->rchild = t->rchild;
    lm->color = t->color;
    rbtree_node_free(t);
    return 0;
}

void rbtree_inorder_print(rbtree_t root)
{
    if (root != NULL)
    {
        rbtree_inorder_print(root->lchild);
        printf("%d<%c>, ", root->data, root->color);
        rbtree_inorder_print(root->rchild);
    }
}

void rbtree_preorder_print(rbtree_t root)
{
    if(root != NULL)
    {
        printf("%d<%c>, ", root->data, root->color);
        rbtree_preorder_print(root->lchild);
        rbtree_preorder_print(root->rchild);
    }
}

void rbtree_postorder_print(rbtree_t root)
{
    if(root != NULL)
    {
        rbtree_postorder_print(root->lchild);
        rbtree_postorder_print(root->rchild);
        printf("%d<%c>, ", root->data, root->color);
    }
}

int main()
{
    int a[] = { 56, 34, 54, 47, 78, 91, 24, 87, 74, 69, 43, 63, 58, 38, 25,91, 49, 23, 12, 8, 76, 85, 102, 99, 29 };
    int num = sizeof(a)/sizeof(int);
    rbtree_t root = NULL;
    for (int i=0; i<num; i++)
    {
        rbtree_insert(&root, a[i]);
    }

    rbtree_inorder_print(root);
    printf("\n");
    rbtree_delete(&root, 35);
    rbtree_inorder_print(root);
    printf("\n");
    rbtree_delete(&root, 77);
    rbtree_inorder_print(root);
}